import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import torch
import torchvision.models as models
import torchvision.transforms as transforms

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei"]

# 模型字典
models_dict = {
    'LeNet': None,
    'VGG16': None,
    'AlexNet': None,
    'GoogLeNet': None
}

# 模型特点字典
model_features = {
    'LeNet': {
        '卷积层特点': '使用5x5卷积核，结构简单，适合小数据集和基础教学',
        '池化层特点': '使用平均池化，逐步缩小特征图尺寸',
        'top1_acc': 0.78,
        'top5_acc': 0.92
    },
    'VGG16': {
        '卷积层特点': '使用多个连续的3x3小卷积核替代大卷积核，减少参数数量，加深网络结构',
        '池化层特点': '固定使用2x2最大池化，步长为2',
        'top1_acc': 0.85,
        'top5_acc': 0.95
    },
    'AlexNet': {
        '卷积层特点': '使用较大卷积核(11x11, 5x5)，首次引入ReLU激活函数，缓解梯度消失问题',
        '池化层特点': '使用重叠最大池化(步长2，窗口3x3)，有助于缓解过拟合',
        'top1_acc': 0.82,
        'top5_acc': 0.93
    },
    'GoogLeNet': {
        '卷积层特点': '引入Inception模块，在同一层使用不同尺寸的卷积核，增加网络宽度',
        '池化层特点': '使用全局平均池化替代全连接层，显著减少参数数量',
        'top1_acc': 0.84,
        'top5_acc': 0.94
    }
}

# 动物类别标签
mnist45_labels = {
    0: "大象",
    1: "羊驼",
    2: "耗牛",
    3: "长嘴小怪兽",
    4: "小狐狸",
    5: "armadillo",
    6: "baboon",
    7: "badger",
    8: "blue_whale",
    9: "brown_bear",
    10: "camel",
    11: "dolphin",
    12: "长颈鹿",
    13: "groundhog",
    14: "highland_cattle",
    15: "马",
    16: "jackal",
    17: "kangaroo",
    18: "koala",
    19: "manatee",
    20: "mongoose",
    21: "mountain_goat",
    22: "opossum",
    23: "orangutan",
    24: "otter",
    25: "polar_bear",
    26: "porcupine",
    27: "小浣熊",
    28: "rhinoceros",
    29: "sea_lion",
    30: "seal",
    31: "snow_leopard",
    32: "squirrel",
    33: "suger_glider",
    34: "tapir",
    35: "vampire_bat",
    36: "vicuna",
    37: "walrus",
    38: "warthog",
    39: "water_buffalo",
    40: "weasel",
    41: "wildebeest",
    42: "wombat",
    43: "yak",
    44: "zebra"
}

# 模型定义
class LeNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 6, 5)
        self.conv2 = torch.nn.Conv2d(6, 16, 5)
        self.conv3 = torch.nn.Conv2d(16, 120, 5)
        self.avg_pool = torch.nn.AvgPool2d(2)
        self.flatten = torch.nn.Flatten(start_dim=1)
        self.tanh = torch.nn.Tanh()
        self.log_softmax = torch.nn.LogSoftmax(dim=1)
        self.f1 = torch.nn.Linear(120, 84)
        self.f2 = torch.nn.Linear(84, 45)  # 输出类别改为45

    def forward(self, x):  # N,3,32,32
        x = self.conv1(x)  # N,6,28,28
        x = self.tanh(x)
        x = self.avg_pool(x)  # N,6,14,14
        x = self.tanh(x)

        x = self.conv2(x)  # N,16,10,10
        x = self.tanh(x)
        x = self.avg_pool(x)  # N,16,5,5
        x = self.tanh(x)

        x = self.conv3(x)  # N,120,1,1
        x = self.tanh(x)
        x = self.flatten(x)  # N,120

        x = self.f1(x)  # N,84
        x = self.tanh(x)
        x = self.f2(x)  # N,45
        y = self.log_softmax(x)

        return y

# 预处理函数
def preprocess_image(image, model_name):
    if model_name == 'LeNet':
        transform = transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
        ])
    else:
        transform = transforms.Compose([
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
    return transform(image)

# 预测函数
def predict_with_model(model, input_tensor, top_k=5):
    model.eval()
    with torch.no_grad():
        outputs = model(input_tensor.unsqueeze(0))
        probs = torch.nn.functional.softmax(outputs, dim=1)
        top_probs, top_indices = torch.topk(probs, top_k, dim=1)
        top_probs = top_probs.squeeze().numpy()
        top_indices = top_indices.squeeze().numpy()
        top_classes = [mnist45_labels[idx] for idx in top_indices]
        predicted_class = top_indices[0]
    return predicted_class, top_classes, top_probs

# 加载模型
def load_models():
    # 加载LeNet
    model_le = LeNet().to('cpu')
    try:
        state_dict = torch.load("static/LeNet_mammals_best_model.pth", map_location=torch.device('cpu'))
        model_le.load_state_dict(state_dict)
    except:
        st.warning("未找到LeNet模型文件，使用随机初始化参数")
    models_dict['LeNet'] = model_le

    # 加载VGG16
    model_vgg = models.vgg16().to('cpu')
    try:
        state_dict = torch.load("static/vgg16_mammals_best_model.pth", map_location=torch.device('cpu'))
        model_vgg.classifier[6] = torch.nn.Linear(4096, 45)  # 修改输出维度为45
        model_vgg.load_state_dict(state_dict)
    except:
        st.warning("未找到VGG16模型文件，使用随机初始化参数")
    models_dict['VGG16'] = model_vgg

    # 加载AlexNet
    model_alex = models.alexnet().to('cpu')
    try:
        state_dict = torch.load("static/alex_mammals_best_model.pth", map_location=torch.device('cpu'))
        model_alex.classifier[6] = torch.nn.Linear(4096, 45)  # 修改输出维度为45
        model_alex.load_state_dict(state_dict)
    except:
        st.warning("未找到AlexNet模型文件，使用随机初始化参数")
    models_dict['AlexNet'] = model_alex

    # 加载GoogLeNet
    model_googlenet = models.googlenet(weights='IMAGENET1K_V1').to('cpu')
    try:
        state_dict = torch.load("static/GoogLeNet_mammals_best_model.pth", map_location=torch.device('cpu'))
        model_googlenet.fc = torch.nn.Linear(1024, 45)  # 修改输出维度为45
        model_googlenet.load_state_dict(state_dict)
    except:
        st.warning("未找到GoogLeNet模型文件，使用随机初始化参数")
    models_dict['GoogLeNet'] = model_googlenet

# 模型对比函数
def compare_models(image):
    load_models()

    results = {}
    for model_name, model in models_dict.items():
        if model is not None:
            input_tensor = preprocess_image(image, model_name)
            predicted_class, top_classes, top_probs = predict_with_model(model, input_tensor)
            results[model_name] = {
                'predicted_class': predicted_class,
                'top_classes': top_classes,
                'top_probs': top_probs
            }

    return results

# 显示对比结果
def display_comparison_results(results, image):
    st.subheader("模型预测结果对比")

    data = []

    for model_name, result in results.items():
        predicted_class_name = mnist45_labels[result['predicted_class']]
        top1_prob = result['top_probs'][0]

        row = {
            '模型': model_name,
            'Top-1预测类别': predicted_class_name,
            'Top-1置信度': f"{top1_prob:.4f}",
            'Top-5预测类别': ", ".join(result['top_classes']),
            'Top-5置信度': ", ".join([f"{prob:.4f}" for prob in result['top_probs']])
        }
        data.append(row)

    # 显示表格
    df = pd.DataFrame(data)
    st.dataframe(df, use_container_width=True)

    # 创建Top-1置信度对比图
    fig, ax = plt.subplots(figsize=(10, 6))
    x = np.arange(len(results))
    top1_probs = [result['top_probs'][0] for result in results.values()]

    ax.bar(x, top1_probs, color='skyblue')
    ax.set_ylabel('置信度')
    ax.set_title('各模型Top-1预测置信度对比')
    ax.set_xticks(x)
    ax.set_xticklabels(list(results.keys()))

    for i, v in enumerate(top1_probs):
        ax.text(i, v + 0.01, f'{v:.4f}', ha='center')

    st.pyplot(fig)

    # 创建模型结构对比表格
    st.subheader("模型结构对比")
    structure_data = []

    for model_name in results.keys():
        structure_data.append({
            '模型': model_name,
            '卷积层特点': model_features[model_name]['卷积层特点'],
            '池化层特点': model_features[model_name]['池化层特点'],
            'ImageNet Top-1准确率': model_features[model_name]['top1_acc'],
            'ImageNet Top-5准确率': model_features[model_name]['top5_acc']
        })

    structure_df = pd.DataFrame(structure_data)
    st.dataframe(structure_df, use_container_width=True)

def main():
    st.title("多模型图片识别系统")
    st.write("请上传一张图片进行识别，支持多轮交互~~~")

    # 初始化会话状态
    if 'history' not in st.session_state:
        st.session_state.history = []

    # 显示历史信息
    for msg in st.session_state.history:
        with st.chat_message(msg['role']):
            if msg['role'] == 'user':
                st.image(msg['content'])
            else:
                st.markdown(msg['content'])

    # 侧边栏设置
    with st.sidebar:
        st.title("侧边栏")
        task = st.selectbox('选择模型', ['LeNet', 'VGG16', 'AlexNet', 'GoogLeNet', '模型对比'])
        st.button("清空历史信息", on_click=lambda: setattr(st.session_state, 'history', []))

    # 关键修改：使用st.file_uploader替代st.chat_input处理图片上传
    uploaded_file = st.file_uploader(
        "上传图片进行识别",
        type=['jpg', 'jpeg', 'png'],
        key="image_uploader"
    )

    # 当有文件上传时处理
    if uploaded_file is not None:
        # 打开图片
        image = Image.open(uploaded_file)

        # 显示用户信息并存储到历史记录
        with st.chat_message("user"):
            st.image(image, caption="用户上传的图片")
        st.session_state.history.append({"role": "user", "content": image})

        if task == "模型对比":
            try:
                results = compare_models(image)
                display_comparison_results(results, image)
            except Exception as e:
                st.error(f"错误: {str(e)}")
                st.session_state.history.append(
                    {"role": "assistant", "content": "识别失败，请检查图片格式或尝试其他图片"})
        else:
            # 单个模型预测
            model_name = task
            model = models_dict.get(model_name)
            if model is None:
                load_models()
                model = models_dict.get(model_name)
            if model is not None:
                try:
                    input_tensor = preprocess_image(image, model_name)
                    predicted_class, top_classes, top_probs = predict_with_model(model, input_tensor)
                    predicted_class_name = mnist45_labels[predicted_class]

                    # 显示系统信息并存储到历史记录
                    with st.chat_message("assistant"):
                        st.markdown(f"**{model_name} 预测结果**")
                        st.markdown(f"Top-1 预测类别: {predicted_class_name}")
                        st.markdown("Top-5 预测结果:")
                        for i, (cls, prob) in enumerate(zip(top_classes, top_probs)):
                            st.markdown(f"{i + 1}. {cls}: {prob:.4f}")

                    st.session_state.history.append(
                        {"role": "assistant", "content": f"{model_name} 预测类别: {predicted_class_name}"})

                    # 绘制Top-5预测结果条形图
                    plt.figure(figsize=(10, 6))
                    plt.barh(top_classes, top_probs)
                    plt.xlabel("预测概率")
                    plt.ylabel("类别名称")
                    plt.title(f"{model_name} Top-5 预测结果")
                    plt.tight_layout()
                    st.pyplot(plt)

                    # 显示模型特点
                    st.markdown(f"**{model_name} 模型特点**")
                    st.markdown(f"卷积层特点: {model_features[model_name]['卷积层特点']}")
                    st.markdown(f"池化层特点: {model_features[model_name]['池化层特点']}")

                except Exception as e:
                    st.error(f"错误: {str(e)}")
                    st.session_state.history.append(
                        {"role": "assistant", "content": "识别失败，请检查图片格式或尝试其他图片"})
            else:
                st.error(f"模型 {model_name} 未找到")
                st.session_state.history.append({"role": "assistant", "content": f"模型 {model_name} 未找到"})

if __name__ == "__main__":
    main()